{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2025-05-27T07:52:27.387675Z",
     "start_time": "2025-05-27T07:52:27.376765Z"
    }
   },
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import Transformer\n",
    "import math\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch.nn.utils.rnn import pad_sequence\n",
    "\n",
    "# 配置参数\n",
    "# DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'mps:0')\n",
    "BATCH_SIZE = 2\n",
    "NUM_EPOCHS = 200\n",
    "# NUM_EPOCHS = 2\n",
    "LEARNING_RATE = 0.001\n",
    "D_MODEL = 128  # 就是EMBEDDING_SIZE\n",
    "NUM_HEAD = 4\n",
    "NUM_ENCODER_LAYERS = 3\n",
    "NUM_DECODER_LAYERS = 3\n",
    "DIM_FEEDFORWARD = 512\n",
    "DROPOUT = 0.1\n",
    "MAX_SEQ_LENGTH = 10  # 推理时用的"
   ],
   "outputs": [],
   "execution_count": 303
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-27T07:52:27.500214Z",
     "start_time": "2025-05-27T07:52:27.492581Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 读取鲁迅文章并预处理\n",
    "with open('luxun_article.txt', 'r', encoding='utf-8') as f:\n",
    "    text = f.read()\n",
    "text"
   ],
   "id": "f1c6fdad671f2635",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'早上，我静坐了一会儿。陈老五送进饭来，一碗菜，一碗蒸鱼；这鱼的眼睛，白而且硬，张着嘴，同那一伙想吃人的人一样。吃了几筷，滑溜溜的不知是鱼是人，便把他兜肚连肠的吐出。我说“老五，对大哥说，我闷得慌，想到园里走走。”老五不答应，走了；停一会，可就来开了门。我也不动，研究他们如何摆布我；知道他们一定不肯放松。果然！我大哥引了一个老头子，慢慢走来；他满眼凶光，怕我看出，只是低头向着地，从眼镜横边暗暗看我。大哥说，“今天你仿佛很好。”我说“是的。”大哥说，“今天请何先生来，给你诊一诊。”我说“可以！”其实我岂不知道这老头子是刽子手扮的！无非借了看脉这名目，揣一揣肥瘠：因这功劳，也分一片肉吃。我也不怕；虽然不吃人，胆子却比他们还壮。伸出两个拳头，看他如何下手。老头子坐着，闭了眼睛，摸了好一会，呆了好一会；便张开他鬼眼睛说，“不要乱想。静静的养几天，就好了。”不要乱想，静静的养！养肥了，他们是自然可以多吃；我有什么好处，怎么会“好了”？他们这群人，又想吃人，又是鬼鬼祟祟，想法子遮掩，不敢直截下手，真要令我笑死。我忍不住，便放声大笑起来，十分快活。自己晓得这笑声里面，有的是义勇和正气。老头子和大哥，都失了色，被我这勇气正气镇压住了。但是我有勇气，他们便越想吃我，沾光一点这勇气。老头子跨出门，走不多远，便低声对大哥说道，“赶紧吃罢！”大哥点点头。原来也有你！这一件大发见，虽似意外，也在意中：合伙吃我的人，便是我的哥哥！吃人的是我哥哥！我是吃人的人的兄弟！我自己被人吃了，可仍然是吃人的人的兄弟！'"
      ]
     },
     "execution_count": 304,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 304
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-27T07:52:27.585418Z",
     "start_time": "2025-05-27T07:52:27.572431Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 数据预处理\n",
    "def preprocess_data(data):\n",
    "    return list(data)\n",
    "\n",
    "\n",
    "tokens = preprocess_data(text)\n",
    "tokens"
   ],
   "id": "7aade81080628def",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['早',\n",
       " '上',\n",
       " '，',\n",
       " '我',\n",
       " '静',\n",
       " '坐',\n",
       " '了',\n",
       " '一',\n",
       " '会',\n",
       " '儿',\n",
       " '。',\n",
       " '陈',\n",
       " '老',\n",
       " '五',\n",
       " '送',\n",
       " '进',\n",
       " '饭',\n",
       " '来',\n",
       " '，',\n",
       " '一',\n",
       " '碗',\n",
       " '菜',\n",
       " '，',\n",
       " '一',\n",
       " '碗',\n",
       " '蒸',\n",
       " '鱼',\n",
       " '；',\n",
       " '这',\n",
       " '鱼',\n",
       " '的',\n",
       " '眼',\n",
       " '睛',\n",
       " '，',\n",
       " '白',\n",
       " '而',\n",
       " '且',\n",
       " '硬',\n",
       " '，',\n",
       " '张',\n",
       " '着',\n",
       " '嘴',\n",
       " '，',\n",
       " '同',\n",
       " '那',\n",
       " '一',\n",
       " '伙',\n",
       " '想',\n",
       " '吃',\n",
       " '人',\n",
       " '的',\n",
       " '人',\n",
       " '一',\n",
       " '样',\n",
       " '。',\n",
       " '吃',\n",
       " '了',\n",
       " '几',\n",
       " '筷',\n",
       " '，',\n",
       " '滑',\n",
       " '溜',\n",
       " '溜',\n",
       " '的',\n",
       " '不',\n",
       " '知',\n",
       " '是',\n",
       " '鱼',\n",
       " '是',\n",
       " '人',\n",
       " '，',\n",
       " '便',\n",
       " '把',\n",
       " '他',\n",
       " '兜',\n",
       " '肚',\n",
       " '连',\n",
       " '肠',\n",
       " '的',\n",
       " '吐',\n",
       " '出',\n",
       " '。',\n",
       " '我',\n",
       " '说',\n",
       " '“',\n",
       " '老',\n",
       " '五',\n",
       " '，',\n",
       " '对',\n",
       " '大',\n",
       " '哥',\n",
       " '说',\n",
       " '，',\n",
       " '我',\n",
       " '闷',\n",
       " '得',\n",
       " '慌',\n",
       " '，',\n",
       " '想',\n",
       " '到',\n",
       " '园',\n",
       " '里',\n",
       " '走',\n",
       " '走',\n",
       " '。',\n",
       " '”',\n",
       " '老',\n",
       " '五',\n",
       " '不',\n",
       " '答',\n",
       " '应',\n",
       " '，',\n",
       " '走',\n",
       " '了',\n",
       " '；',\n",
       " '停',\n",
       " '一',\n",
       " '会',\n",
       " '，',\n",
       " '可',\n",
       " '就',\n",
       " '来',\n",
       " '开',\n",
       " '了',\n",
       " '门',\n",
       " '。',\n",
       " '我',\n",
       " '也',\n",
       " '不',\n",
       " '动',\n",
       " '，',\n",
       " '研',\n",
       " '究',\n",
       " '他',\n",
       " '们',\n",
       " '如',\n",
       " '何',\n",
       " '摆',\n",
       " '布',\n",
       " '我',\n",
       " '；',\n",
       " '知',\n",
       " '道',\n",
       " '他',\n",
       " '们',\n",
       " '一',\n",
       " '定',\n",
       " '不',\n",
       " '肯',\n",
       " '放',\n",
       " '松',\n",
       " '。',\n",
       " '果',\n",
       " '然',\n",
       " '！',\n",
       " '我',\n",
       " '大',\n",
       " '哥',\n",
       " '引',\n",
       " '了',\n",
       " '一',\n",
       " '个',\n",
       " '老',\n",
       " '头',\n",
       " '子',\n",
       " '，',\n",
       " '慢',\n",
       " '慢',\n",
       " '走',\n",
       " '来',\n",
       " '；',\n",
       " '他',\n",
       " '满',\n",
       " '眼',\n",
       " '凶',\n",
       " '光',\n",
       " '，',\n",
       " '怕',\n",
       " '我',\n",
       " '看',\n",
       " '出',\n",
       " '，',\n",
       " '只',\n",
       " '是',\n",
       " '低',\n",
       " '头',\n",
       " '向',\n",
       " '着',\n",
       " '地',\n",
       " '，',\n",
       " '从',\n",
       " '眼',\n",
       " '镜',\n",
       " '横',\n",
       " '边',\n",
       " '暗',\n",
       " '暗',\n",
       " '看',\n",
       " '我',\n",
       " '。',\n",
       " '大',\n",
       " '哥',\n",
       " '说',\n",
       " '，',\n",
       " '“',\n",
       " '今',\n",
       " '天',\n",
       " '你',\n",
       " '仿',\n",
       " '佛',\n",
       " '很',\n",
       " '好',\n",
       " '。',\n",
       " '”',\n",
       " '我',\n",
       " '说',\n",
       " '“',\n",
       " '是',\n",
       " '的',\n",
       " '。',\n",
       " '”',\n",
       " '大',\n",
       " '哥',\n",
       " '说',\n",
       " '，',\n",
       " '“',\n",
       " '今',\n",
       " '天',\n",
       " '请',\n",
       " '何',\n",
       " '先',\n",
       " '生',\n",
       " '来',\n",
       " '，',\n",
       " '给',\n",
       " '你',\n",
       " '诊',\n",
       " '一',\n",
       " '诊',\n",
       " '。',\n",
       " '”',\n",
       " '我',\n",
       " '说',\n",
       " '“',\n",
       " '可',\n",
       " '以',\n",
       " '！',\n",
       " '”',\n",
       " '其',\n",
       " '实',\n",
       " '我',\n",
       " '岂',\n",
       " '不',\n",
       " '知',\n",
       " '道',\n",
       " '这',\n",
       " '老',\n",
       " '头',\n",
       " '子',\n",
       " '是',\n",
       " '刽',\n",
       " '子',\n",
       " '手',\n",
       " '扮',\n",
       " '的',\n",
       " '！',\n",
       " '无',\n",
       " '非',\n",
       " '借',\n",
       " '了',\n",
       " '看',\n",
       " '脉',\n",
       " '这',\n",
       " '名',\n",
       " '目',\n",
       " '，',\n",
       " '揣',\n",
       " '一',\n",
       " '揣',\n",
       " '肥',\n",
       " '瘠',\n",
       " '：',\n",
       " '因',\n",
       " '这',\n",
       " '功',\n",
       " '劳',\n",
       " '，',\n",
       " '也',\n",
       " '分',\n",
       " '一',\n",
       " '片',\n",
       " '肉',\n",
       " '吃',\n",
       " '。',\n",
       " '我',\n",
       " '也',\n",
       " '不',\n",
       " '怕',\n",
       " '；',\n",
       " '虽',\n",
       " '然',\n",
       " '不',\n",
       " '吃',\n",
       " '人',\n",
       " '，',\n",
       " '胆',\n",
       " '子',\n",
       " '却',\n",
       " '比',\n",
       " '他',\n",
       " '们',\n",
       " '还',\n",
       " '壮',\n",
       " '。',\n",
       " '伸',\n",
       " '出',\n",
       " '两',\n",
       " '个',\n",
       " '拳',\n",
       " '头',\n",
       " '，',\n",
       " '看',\n",
       " '他',\n",
       " '如',\n",
       " '何',\n",
       " '下',\n",
       " '手',\n",
       " '。',\n",
       " '老',\n",
       " '头',\n",
       " '子',\n",
       " '坐',\n",
       " '着',\n",
       " '，',\n",
       " '闭',\n",
       " '了',\n",
       " '眼',\n",
       " '睛',\n",
       " '，',\n",
       " '摸',\n",
       " '了',\n",
       " '好',\n",
       " '一',\n",
       " '会',\n",
       " '，',\n",
       " '呆',\n",
       " '了',\n",
       " '好',\n",
       " '一',\n",
       " '会',\n",
       " '；',\n",
       " '便',\n",
       " '张',\n",
       " '开',\n",
       " '他',\n",
       " '鬼',\n",
       " '眼',\n",
       " '睛',\n",
       " '说',\n",
       " '，',\n",
       " '“',\n",
       " '不',\n",
       " '要',\n",
       " '乱',\n",
       " '想',\n",
       " '。',\n",
       " '静',\n",
       " '静',\n",
       " '的',\n",
       " '养',\n",
       " '几',\n",
       " '天',\n",
       " '，',\n",
       " '就',\n",
       " '好',\n",
       " '了',\n",
       " '。',\n",
       " '”',\n",
       " '不',\n",
       " '要',\n",
       " '乱',\n",
       " '想',\n",
       " '，',\n",
       " '静',\n",
       " '静',\n",
       " '的',\n",
       " '养',\n",
       " '！',\n",
       " '养',\n",
       " '肥',\n",
       " '了',\n",
       " '，',\n",
       " '他',\n",
       " '们',\n",
       " '是',\n",
       " '自',\n",
       " '然',\n",
       " '可',\n",
       " '以',\n",
       " '多',\n",
       " '吃',\n",
       " '；',\n",
       " '我',\n",
       " '有',\n",
       " '什',\n",
       " '么',\n",
       " '好',\n",
       " '处',\n",
       " '，',\n",
       " '怎',\n",
       " '么',\n",
       " '会',\n",
       " '“',\n",
       " '好',\n",
       " '了',\n",
       " '”',\n",
       " '？',\n",
       " '他',\n",
       " '们',\n",
       " '这',\n",
       " '群',\n",
       " '人',\n",
       " '，',\n",
       " '又',\n",
       " '想',\n",
       " '吃',\n",
       " '人',\n",
       " '，',\n",
       " '又',\n",
       " '是',\n",
       " '鬼',\n",
       " '鬼',\n",
       " '祟',\n",
       " '祟',\n",
       " '，',\n",
       " '想',\n",
       " '法',\n",
       " '子',\n",
       " '遮',\n",
       " '掩',\n",
       " '，',\n",
       " '不',\n",
       " '敢',\n",
       " '直',\n",
       " '截',\n",
       " '下',\n",
       " '手',\n",
       " '，',\n",
       " '真',\n",
       " '要',\n",
       " '令',\n",
       " '我',\n",
       " '笑',\n",
       " '死',\n",
       " '。',\n",
       " '我',\n",
       " '忍',\n",
       " '不',\n",
       " '住',\n",
       " '，',\n",
       " '便',\n",
       " '放',\n",
       " '声',\n",
       " '大',\n",
       " '笑',\n",
       " '起',\n",
       " '来',\n",
       " '，',\n",
       " '十',\n",
       " '分',\n",
       " '快',\n",
       " '活',\n",
       " '。',\n",
       " '自',\n",
       " '己',\n",
       " '晓',\n",
       " '得',\n",
       " '这',\n",
       " '笑',\n",
       " '声',\n",
       " '里',\n",
       " '面',\n",
       " '，',\n",
       " '有',\n",
       " '的',\n",
       " '是',\n",
       " '义',\n",
       " '勇',\n",
       " '和',\n",
       " '正',\n",
       " '气',\n",
       " '。',\n",
       " '老',\n",
       " '头',\n",
       " '子',\n",
       " '和',\n",
       " '大',\n",
       " '哥',\n",
       " '，',\n",
       " '都',\n",
       " '失',\n",
       " '了',\n",
       " '色',\n",
       " '，',\n",
       " '被',\n",
       " '我',\n",
       " '这',\n",
       " '勇',\n",
       " '气',\n",
       " '正',\n",
       " '气',\n",
       " '镇',\n",
       " '压',\n",
       " '住',\n",
       " '了',\n",
       " '。',\n",
       " '但',\n",
       " '是',\n",
       " '我',\n",
       " '有',\n",
       " '勇',\n",
       " '气',\n",
       " '，',\n",
       " '他',\n",
       " '们',\n",
       " '便',\n",
       " '越',\n",
       " '想',\n",
       " '吃',\n",
       " '我',\n",
       " '，',\n",
       " '沾',\n",
       " '光',\n",
       " '一',\n",
       " '点',\n",
       " '这',\n",
       " '勇',\n",
       " '气',\n",
       " '。',\n",
       " '老',\n",
       " '头',\n",
       " '子',\n",
       " '跨',\n",
       " '出',\n",
       " '门',\n",
       " '，',\n",
       " '走',\n",
       " '不',\n",
       " '多',\n",
       " '远',\n",
       " '，',\n",
       " '便',\n",
       " '低',\n",
       " '声',\n",
       " '对',\n",
       " '大',\n",
       " '哥',\n",
       " '说',\n",
       " '道',\n",
       " '，',\n",
       " '“',\n",
       " '赶',\n",
       " '紧',\n",
       " '吃',\n",
       " '罢',\n",
       " '！',\n",
       " '”',\n",
       " '大',\n",
       " '哥',\n",
       " '点',\n",
       " '点',\n",
       " '头',\n",
       " '。',\n",
       " '原',\n",
       " '来',\n",
       " '也',\n",
       " '有',\n",
       " '你',\n",
       " '！',\n",
       " '这',\n",
       " '一',\n",
       " '件',\n",
       " '大',\n",
       " '发',\n",
       " '见',\n",
       " '，',\n",
       " '虽',\n",
       " '似',\n",
       " '意',\n",
       " '外',\n",
       " '，',\n",
       " '也',\n",
       " '在',\n",
       " '意',\n",
       " '中',\n",
       " '：',\n",
       " '合',\n",
       " '伙',\n",
       " '吃',\n",
       " '我',\n",
       " '的',\n",
       " '人',\n",
       " '，',\n",
       " '便',\n",
       " '是',\n",
       " '我',\n",
       " '的',\n",
       " '哥',\n",
       " '哥',\n",
       " '！',\n",
       " '吃',\n",
       " '人',\n",
       " '的',\n",
       " '是',\n",
       " '我',\n",
       " '哥',\n",
       " '哥',\n",
       " '！',\n",
       " '我',\n",
       " '是',\n",
       " '吃',\n",
       " '人',\n",
       " '的',\n",
       " '人',\n",
       " '的',\n",
       " '兄',\n",
       " '弟',\n",
       " '！',\n",
       " '我',\n",
       " '自',\n",
       " '己',\n",
       " '被',\n",
       " '人',\n",
       " '吃',\n",
       " '了',\n",
       " '，',\n",
       " '可',\n",
       " '仍',\n",
       " '然',\n",
       " '是',\n",
       " '吃',\n",
       " '人',\n",
       " '的',\n",
       " '人',\n",
       " '的',\n",
       " '兄',\n",
       " '弟',\n",
       " '！']"
      ]
     },
     "execution_count": 305,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 305
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-27T07:52:27.632119Z",
     "start_time": "2025-05-27T07:52:27.626627Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 构建词汇表\n",
    "def build_vocab(tokens):\n",
    "    vocab = {'<pad>': 0, '<sos>': 1, '<eos>': 2, '<unk>': 3}\n",
    "    for token in tokens:\n",
    "        if token not in vocab:\n",
    "            vocab[token] = len(vocab)\n",
    "    return vocab\n",
    "\n",
    "\n",
    "vocab = build_vocab(tokens)\n",
    "\n",
    "vocab"
   ],
   "id": "cb470ec639e0c102",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'<pad>': 0,\n",
       " '<sos>': 1,\n",
       " '<eos>': 2,\n",
       " '<unk>': 3,\n",
       " '早': 4,\n",
       " '上': 5,\n",
       " '，': 6,\n",
       " '我': 7,\n",
       " '静': 8,\n",
       " '坐': 9,\n",
       " '了': 10,\n",
       " '一': 11,\n",
       " '会': 12,\n",
       " '儿': 13,\n",
       " '。': 14,\n",
       " '陈': 15,\n",
       " '老': 16,\n",
       " '五': 17,\n",
       " '送': 18,\n",
       " '进': 19,\n",
       " '饭': 20,\n",
       " '来': 21,\n",
       " '碗': 22,\n",
       " '菜': 23,\n",
       " '蒸': 24,\n",
       " '鱼': 25,\n",
       " '；': 26,\n",
       " '这': 27,\n",
       " '的': 28,\n",
       " '眼': 29,\n",
       " '睛': 30,\n",
       " '白': 31,\n",
       " '而': 32,\n",
       " '且': 33,\n",
       " '硬': 34,\n",
       " '张': 35,\n",
       " '着': 36,\n",
       " '嘴': 37,\n",
       " '同': 38,\n",
       " '那': 39,\n",
       " '伙': 40,\n",
       " '想': 41,\n",
       " '吃': 42,\n",
       " '人': 43,\n",
       " '样': 44,\n",
       " '几': 45,\n",
       " '筷': 46,\n",
       " '滑': 47,\n",
       " '溜': 48,\n",
       " '不': 49,\n",
       " '知': 50,\n",
       " '是': 51,\n",
       " '便': 52,\n",
       " '把': 53,\n",
       " '他': 54,\n",
       " '兜': 55,\n",
       " '肚': 56,\n",
       " '连': 57,\n",
       " '肠': 58,\n",
       " '吐': 59,\n",
       " '出': 60,\n",
       " '说': 61,\n",
       " '“': 62,\n",
       " '对': 63,\n",
       " '大': 64,\n",
       " '哥': 65,\n",
       " '闷': 66,\n",
       " '得': 67,\n",
       " '慌': 68,\n",
       " '到': 69,\n",
       " '园': 70,\n",
       " '里': 71,\n",
       " '走': 72,\n",
       " '”': 73,\n",
       " '答': 74,\n",
       " '应': 75,\n",
       " '停': 76,\n",
       " '可': 77,\n",
       " '就': 78,\n",
       " '开': 79,\n",
       " '门': 80,\n",
       " '也': 81,\n",
       " '动': 82,\n",
       " '研': 83,\n",
       " '究': 84,\n",
       " '们': 85,\n",
       " '如': 86,\n",
       " '何': 87,\n",
       " '摆': 88,\n",
       " '布': 89,\n",
       " '道': 90,\n",
       " '定': 91,\n",
       " '肯': 92,\n",
       " '放': 93,\n",
       " '松': 94,\n",
       " '果': 95,\n",
       " '然': 96,\n",
       " '！': 97,\n",
       " '引': 98,\n",
       " '个': 99,\n",
       " '头': 100,\n",
       " '子': 101,\n",
       " '慢': 102,\n",
       " '满': 103,\n",
       " '凶': 104,\n",
       " '光': 105,\n",
       " '怕': 106,\n",
       " '看': 107,\n",
       " '只': 108,\n",
       " '低': 109,\n",
       " '向': 110,\n",
       " '地': 111,\n",
       " '从': 112,\n",
       " '镜': 113,\n",
       " '横': 114,\n",
       " '边': 115,\n",
       " '暗': 116,\n",
       " '今': 117,\n",
       " '天': 118,\n",
       " '你': 119,\n",
       " '仿': 120,\n",
       " '佛': 121,\n",
       " '很': 122,\n",
       " '好': 123,\n",
       " '请': 124,\n",
       " '先': 125,\n",
       " '生': 126,\n",
       " '给': 127,\n",
       " '诊': 128,\n",
       " '以': 129,\n",
       " '其': 130,\n",
       " '实': 131,\n",
       " '岂': 132,\n",
       " '刽': 133,\n",
       " '手': 134,\n",
       " '扮': 135,\n",
       " '无': 136,\n",
       " '非': 137,\n",
       " '借': 138,\n",
       " '脉': 139,\n",
       " '名': 140,\n",
       " '目': 141,\n",
       " '揣': 142,\n",
       " '肥': 143,\n",
       " '瘠': 144,\n",
       " '：': 145,\n",
       " '因': 146,\n",
       " '功': 147,\n",
       " '劳': 148,\n",
       " '分': 149,\n",
       " '片': 150,\n",
       " '肉': 151,\n",
       " '虽': 152,\n",
       " '胆': 153,\n",
       " '却': 154,\n",
       " '比': 155,\n",
       " '还': 156,\n",
       " '壮': 157,\n",
       " '伸': 158,\n",
       " '两': 159,\n",
       " '拳': 160,\n",
       " '下': 161,\n",
       " '闭': 162,\n",
       " '摸': 163,\n",
       " '呆': 164,\n",
       " '鬼': 165,\n",
       " '要': 166,\n",
       " '乱': 167,\n",
       " '养': 168,\n",
       " '自': 169,\n",
       " '多': 170,\n",
       " '有': 171,\n",
       " '什': 172,\n",
       " '么': 173,\n",
       " '处': 174,\n",
       " '怎': 175,\n",
       " '？': 176,\n",
       " '群': 177,\n",
       " '又': 178,\n",
       " '祟': 179,\n",
       " '法': 180,\n",
       " '遮': 181,\n",
       " '掩': 182,\n",
       " '敢': 183,\n",
       " '直': 184,\n",
       " '截': 185,\n",
       " '真': 186,\n",
       " '令': 187,\n",
       " '笑': 188,\n",
       " '死': 189,\n",
       " '忍': 190,\n",
       " '住': 191,\n",
       " '声': 192,\n",
       " '起': 193,\n",
       " '十': 194,\n",
       " '快': 195,\n",
       " '活': 196,\n",
       " '己': 197,\n",
       " '晓': 198,\n",
       " '面': 199,\n",
       " '义': 200,\n",
       " '勇': 201,\n",
       " '和': 202,\n",
       " '正': 203,\n",
       " '气': 204,\n",
       " '都': 205,\n",
       " '失': 206,\n",
       " '色': 207,\n",
       " '被': 208,\n",
       " '镇': 209,\n",
       " '压': 210,\n",
       " '但': 211,\n",
       " '越': 212,\n",
       " '沾': 213,\n",
       " '点': 214,\n",
       " '跨': 215,\n",
       " '远': 216,\n",
       " '赶': 217,\n",
       " '紧': 218,\n",
       " '罢': 219,\n",
       " '原': 220,\n",
       " '件': 221,\n",
       " '发': 222,\n",
       " '见': 223,\n",
       " '似': 224,\n",
       " '意': 225,\n",
       " '外': 226,\n",
       " '在': 227,\n",
       " '中': 228,\n",
       " '合': 229,\n",
       " '兄': 230,\n",
       " '弟': 231,\n",
       " '仍': 232}"
      ]
     },
     "execution_count": 306,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 306
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-27T07:52:27.678562Z",
     "start_time": "2025-05-27T07:52:27.673583Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 反转词汇表用于解码\n",
    "idx_to_tgt = {v: k for k, v in vocab.items()}\n",
    "idx_to_tgt"
   ],
   "id": "87fa6708baf80ca4",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{0: '<pad>',\n",
       " 1: '<sos>',\n",
       " 2: '<eos>',\n",
       " 3: '<unk>',\n",
       " 4: '早',\n",
       " 5: '上',\n",
       " 6: '，',\n",
       " 7: '我',\n",
       " 8: '静',\n",
       " 9: '坐',\n",
       " 10: '了',\n",
       " 11: '一',\n",
       " 12: '会',\n",
       " 13: '儿',\n",
       " 14: '。',\n",
       " 15: '陈',\n",
       " 16: '老',\n",
       " 17: '五',\n",
       " 18: '送',\n",
       " 19: '进',\n",
       " 20: '饭',\n",
       " 21: '来',\n",
       " 22: '碗',\n",
       " 23: '菜',\n",
       " 24: '蒸',\n",
       " 25: '鱼',\n",
       " 26: '；',\n",
       " 27: '这',\n",
       " 28: '的',\n",
       " 29: '眼',\n",
       " 30: '睛',\n",
       " 31: '白',\n",
       " 32: '而',\n",
       " 33: '且',\n",
       " 34: '硬',\n",
       " 35: '张',\n",
       " 36: '着',\n",
       " 37: '嘴',\n",
       " 38: '同',\n",
       " 39: '那',\n",
       " 40: '伙',\n",
       " 41: '想',\n",
       " 42: '吃',\n",
       " 43: '人',\n",
       " 44: '样',\n",
       " 45: '几',\n",
       " 46: '筷',\n",
       " 47: '滑',\n",
       " 48: '溜',\n",
       " 49: '不',\n",
       " 50: '知',\n",
       " 51: '是',\n",
       " 52: '便',\n",
       " 53: '把',\n",
       " 54: '他',\n",
       " 55: '兜',\n",
       " 56: '肚',\n",
       " 57: '连',\n",
       " 58: '肠',\n",
       " 59: '吐',\n",
       " 60: '出',\n",
       " 61: '说',\n",
       " 62: '“',\n",
       " 63: '对',\n",
       " 64: '大',\n",
       " 65: '哥',\n",
       " 66: '闷',\n",
       " 67: '得',\n",
       " 68: '慌',\n",
       " 69: '到',\n",
       " 70: '园',\n",
       " 71: '里',\n",
       " 72: '走',\n",
       " 73: '”',\n",
       " 74: '答',\n",
       " 75: '应',\n",
       " 76: '停',\n",
       " 77: '可',\n",
       " 78: '就',\n",
       " 79: '开',\n",
       " 80: '门',\n",
       " 81: '也',\n",
       " 82: '动',\n",
       " 83: '研',\n",
       " 84: '究',\n",
       " 85: '们',\n",
       " 86: '如',\n",
       " 87: '何',\n",
       " 88: '摆',\n",
       " 89: '布',\n",
       " 90: '道',\n",
       " 91: '定',\n",
       " 92: '肯',\n",
       " 93: '放',\n",
       " 94: '松',\n",
       " 95: '果',\n",
       " 96: '然',\n",
       " 97: '！',\n",
       " 98: '引',\n",
       " 99: '个',\n",
       " 100: '头',\n",
       " 101: '子',\n",
       " 102: '慢',\n",
       " 103: '满',\n",
       " 104: '凶',\n",
       " 105: '光',\n",
       " 106: '怕',\n",
       " 107: '看',\n",
       " 108: '只',\n",
       " 109: '低',\n",
       " 110: '向',\n",
       " 111: '地',\n",
       " 112: '从',\n",
       " 113: '镜',\n",
       " 114: '横',\n",
       " 115: '边',\n",
       " 116: '暗',\n",
       " 117: '今',\n",
       " 118: '天',\n",
       " 119: '你',\n",
       " 120: '仿',\n",
       " 121: '佛',\n",
       " 122: '很',\n",
       " 123: '好',\n",
       " 124: '请',\n",
       " 125: '先',\n",
       " 126: '生',\n",
       " 127: '给',\n",
       " 128: '诊',\n",
       " 129: '以',\n",
       " 130: '其',\n",
       " 131: '实',\n",
       " 132: '岂',\n",
       " 133: '刽',\n",
       " 134: '手',\n",
       " 135: '扮',\n",
       " 136: '无',\n",
       " 137: '非',\n",
       " 138: '借',\n",
       " 139: '脉',\n",
       " 140: '名',\n",
       " 141: '目',\n",
       " 142: '揣',\n",
       " 143: '肥',\n",
       " 144: '瘠',\n",
       " 145: '：',\n",
       " 146: '因',\n",
       " 147: '功',\n",
       " 148: '劳',\n",
       " 149: '分',\n",
       " 150: '片',\n",
       " 151: '肉',\n",
       " 152: '虽',\n",
       " 153: '胆',\n",
       " 154: '却',\n",
       " 155: '比',\n",
       " 156: '还',\n",
       " 157: '壮',\n",
       " 158: '伸',\n",
       " 159: '两',\n",
       " 160: '拳',\n",
       " 161: '下',\n",
       " 162: '闭',\n",
       " 163: '摸',\n",
       " 164: '呆',\n",
       " 165: '鬼',\n",
       " 166: '要',\n",
       " 167: '乱',\n",
       " 168: '养',\n",
       " 169: '自',\n",
       " 170: '多',\n",
       " 171: '有',\n",
       " 172: '什',\n",
       " 173: '么',\n",
       " 174: '处',\n",
       " 175: '怎',\n",
       " 176: '？',\n",
       " 177: '群',\n",
       " 178: '又',\n",
       " 179: '祟',\n",
       " 180: '法',\n",
       " 181: '遮',\n",
       " 182: '掩',\n",
       " 183: '敢',\n",
       " 184: '直',\n",
       " 185: '截',\n",
       " 186: '真',\n",
       " 187: '令',\n",
       " 188: '笑',\n",
       " 189: '死',\n",
       " 190: '忍',\n",
       " 191: '住',\n",
       " 192: '声',\n",
       " 193: '起',\n",
       " 194: '十',\n",
       " 195: '快',\n",
       " 196: '活',\n",
       " 197: '己',\n",
       " 198: '晓',\n",
       " 199: '面',\n",
       " 200: '义',\n",
       " 201: '勇',\n",
       " 202: '和',\n",
       " 203: '正',\n",
       " 204: '气',\n",
       " 205: '都',\n",
       " 206: '失',\n",
       " 207: '色',\n",
       " 208: '被',\n",
       " 209: '镇',\n",
       " 210: '压',\n",
       " 211: '但',\n",
       " 212: '越',\n",
       " 213: '沾',\n",
       " 214: '点',\n",
       " 215: '跨',\n",
       " 216: '远',\n",
       " 217: '赶',\n",
       " 218: '紧',\n",
       " 219: '罢',\n",
       " 220: '原',\n",
       " 221: '件',\n",
       " 222: '发',\n",
       " 223: '见',\n",
       " 224: '似',\n",
       " 225: '意',\n",
       " 226: '外',\n",
       " 227: '在',\n",
       " 228: '中',\n",
       " 229: '合',\n",
       " 230: '兄',\n",
       " 231: '弟',\n",
       " 232: '仍'}"
      ]
     },
     "execution_count": 307,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 307
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-27T07:52:27.722192Z",
     "start_time": "2025-05-27T07:52:27.716253Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 自定义Dataset\n",
    "class GPTDataset(Dataset):\n",
    "    def __init__(self, tokens, vocab, window_size=5):\n",
    "        self.data = []\n",
    "        for i in range(0, len(tokens) - window_size, window_size):\n",
    "            self.data.append((\n",
    "                torch.tensor([vocab[token] for token in tokens[i:i + window_size]], dtype=torch.long),\n",
    "                torch.tensor([vocab[token] for token in tokens[i + 1:i + window_size + 1]], dtype=torch.long),\n",
    "            ))\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.data[idx]\n",
    "\n",
    "\n",
    "dataset = GPTDataset(tokens, vocab)\n",
    "dataset[-2:]"
   ],
   "id": "9d1c27155a783a6b",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(tensor([  6,  77, 232,  96,  51]), tensor([ 77, 232,  96,  51,  42])),\n",
       " (tensor([42, 43, 28, 43, 28]), tensor([ 43,  28,  43,  28, 230]))]"
      ]
     },
     "execution_count": 308,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 308
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-27T07:52:27.767175Z",
     "start_time": "2025-05-27T07:52:27.762850Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 数据加载器\n",
    "data_loader = DataLoader(dataset, batch_size=BATCH_SIZE)"
   ],
   "id": "26f6e536375b476b",
   "outputs": [],
   "execution_count": 309
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-27T07:52:27.816274Z",
     "start_time": "2025-05-27T07:52:27.811454Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 位置编码\n",
    "# 相当于生成了一个很大的位置编码，根据x的seq_len直接获取对应长度的位置编码\n",
    "class PositionalEncoding(nn.Module):\n",
    "    def __init__(self, d_model, max_len=5000):\n",
    "        super().__init__()\n",
    "        pe = torch.zeros(max_len, d_model)\n",
    "        pos = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)\n",
    "        pe[:, 0::2] = torch.sin(pos / (10000 ** (2 * torch.arange(0, d_model, 2) / d_model)))\n",
    "        pe[:, 1::2] = torch.cos(pos / (10000 ** (2 * torch.arange(1, d_model, 2) / d_model)))\n",
    "        self.pe = pe.unsqueeze(0).to(DEVICE)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x + self.pe[:, :x.size(1), :]\n",
    "        return x"
   ],
   "id": "b4c6148adffb1b54",
   "outputs": [],
   "execution_count": 310
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-27T07:52:27.859280Z",
     "start_time": "2025-05-27T07:52:27.854943Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# Feed Forward，两个线程层，最后输出维度不变\n",
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, embed_dim, fc_dim):\n",
    "        super().__init__()\n",
    "        self.fc1 = nn.Linear(embed_dim, fc_dim)\n",
    "        self.fc2 = nn.Linear(fc_dim, embed_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.fc2(torch.relu(self.fc1(x)))\n",
    "\n",
    "# 定义Decoder\n",
    "class Decoder(nn.Module):\n",
    "    def __init__(self, embed_dim, fc_dim, num_heads, num_layers):\n",
    "        super().__init__()\n",
    "        self.num_layers = num_layers\n",
    "        self.num_heads = num_heads\n",
    "        self.layers = nn.ModuleList([\n",
    "            DecoderLayer(embed_dim, fc_dim, num_heads)\n",
    "            for _ in range(num_layers)\n",
    "        ])\n",
    "\n",
    "\n",
    "    def forward(self, inputs_decoder, mask=None):\n",
    "        for layer in self.layers:\n",
    "            inputs_decoder = layer(inputs_decoder, mask=mask)\n",
    "\n",
    "        return inputs_decoder\n",
    "\n",
    "\n",
    "# 定义DecoderLayer\n",
    "class DecoderLayer(nn.Module):\n",
    "    def __init__(self, embed_dim, fc_dim, num_heads):\n",
    "        super().__init__()\n",
    "        self.mha = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)\n",
    "        self.feed_forward = FeedForward(embed_dim, fc_dim)\n",
    "        self.layer_norm1 = nn.LayerNorm(embed_dim)\n",
    "        self.layer_norm2 = nn.LayerNorm(embed_dim)\n",
    "\n",
    "    def forward(self, inputs_decoder, mask=None):\n",
    "        attn_output, _ = self.mha(query=inputs_decoder, key=inputs_decoder, value=inputs_decoder, attn_mask=mask)\n",
    "\n",
    "        outputs_attention = self.layer_norm1(inputs_decoder + attn_output)\n",
    "\n",
    "        return self.layer_norm2(outputs_attention + self.feed_forward(outputs_attention))"
   ],
   "id": "a534618e13744964",
   "outputs": [],
   "execution_count": 311
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-27T07:52:27.900187Z",
     "start_time": "2025-05-27T07:52:27.897271Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# GPT模型\n",
    "class GPTModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.embedding = nn.Embedding(len(vocab), D_MODEL)\n",
    "        self.pos_encoder = PositionalEncoding(D_MODEL)\n",
    "        self.decoder = Decoder(\n",
    "            embed_dim=D_MODEL,\n",
    "            num_heads=NUM_HEAD,\n",
    "            num_layers=NUM_DECODER_LAYERS,\n",
    "            fc_dim=D_MODEL*4\n",
    "        )\n",
    "        self.linear = nn.Linear(D_MODEL, len(vocab))\n",
    "\n",
    "    def forward(self, src, mask=None):\n",
    "        # 嵌入和位置编码\n",
    "        src_emb = self.pos_encoder(self.embedding(src))\n",
    "\n",
    "        # 仅使用解码器模式\n",
    "        output = self.decoder(\n",
    "            src_emb,\n",
    "            mask=mask,\n",
    "        )\n",
    "\n",
    "        return self.linear(output)"
   ],
   "id": "8e52097d64495e52",
   "outputs": [],
   "execution_count": 312
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-27T07:54:58.686759Z",
     "start_time": "2025-05-27T07:52:27.940215Z"
    }
   },
   "cell_type": "code",
   "source": [
    "model = GPTModel().to(DEVICE)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "# 训练循环\n",
    "for epoch in range(NUM_EPOCHS):\n",
    "    model.train()\n",
    "    total_loss = 0\n",
    "\n",
    "    for src, tgt in data_loader:\n",
    "        src = src.to(DEVICE)\n",
    "        tgt = tgt.to(DEVICE)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # 生成因果掩码 (shape: [batch_size, num_heads, seq_len, seq_len])\n",
    "        batch_size, seq_len = src.size(0), src.size(1)\n",
    "        mask = torch.tril(torch.ones(seq_len, seq_len)).to(DEVICE)\n",
    "\n",
    "        output = model(src, mask=mask)\n",
    "\n",
    "        # 调整输出形状\n",
    "        loss = criterion(output.view(-1, output.size(-1)), tgt.view(-1))\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        total_loss += loss.item()\n",
    "\n",
    "    avg_loss = total_loss / len(data_loader)\n",
    "    print(f'Epoch [{epoch + 1}/{NUM_EPOCHS}], Loss: {avg_loss:.4f}')"
   ],
   "id": "2eb65b25b022f95f",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/200], Loss: 5.3029\n",
      "Epoch [2/200], Loss: 4.3207\n",
      "Epoch [3/200], Loss: 3.0868\n",
      "Epoch [4/200], Loss: 2.1707\n",
      "Epoch [5/200], Loss: 1.3957\n",
      "Epoch [6/200], Loss: 0.7940\n",
      "Epoch [7/200], Loss: 0.4273\n",
      "Epoch [8/200], Loss: 0.2834\n",
      "Epoch [9/200], Loss: 0.2078\n",
      "Epoch [10/200], Loss: 0.2042\n",
      "Epoch [11/200], Loss: 0.1912\n",
      "Epoch [12/200], Loss: 0.1482\n",
      "Epoch [13/200], Loss: 0.1229\n",
      "Epoch [14/200], Loss: 0.1085\n",
      "Epoch [15/200], Loss: 0.0945\n",
      "Epoch [16/200], Loss: 0.1023\n",
      "Epoch [17/200], Loss: 0.0793\n",
      "Epoch [18/200], Loss: 0.0797\n",
      "Epoch [19/200], Loss: 0.0753\n",
      "Epoch [20/200], Loss: 0.0732\n",
      "Epoch [21/200], Loss: 0.0659\n",
      "Epoch [22/200], Loss: 0.0726\n",
      "Epoch [23/200], Loss: 0.0758\n",
      "Epoch [24/200], Loss: 0.0748\n",
      "Epoch [25/200], Loss: 0.0751\n",
      "Epoch [26/200], Loss: 0.0731\n",
      "Epoch [27/200], Loss: 0.0758\n",
      "Epoch [28/200], Loss: 0.1618\n",
      "Epoch [29/200], Loss: 0.6486\n",
      "Epoch [30/200], Loss: 1.2545\n",
      "Epoch [31/200], Loss: 0.9606\n",
      "Epoch [32/200], Loss: 0.4488\n",
      "Epoch [33/200], Loss: 0.2229\n",
      "Epoch [34/200], Loss: 0.1485\n",
      "Epoch [35/200], Loss: 0.1442\n",
      "Epoch [36/200], Loss: 0.1232\n",
      "Epoch [37/200], Loss: 0.1047\n",
      "Epoch [38/200], Loss: 0.0895\n",
      "Epoch [39/200], Loss: 0.0866\n",
      "Epoch [40/200], Loss: 0.0680\n",
      "Epoch [41/200], Loss: 0.0673\n",
      "Epoch [42/200], Loss: 0.0646\n",
      "Epoch [43/200], Loss: 0.0647\n",
      "Epoch [44/200], Loss: 0.0548\n",
      "Epoch [45/200], Loss: 0.0597\n",
      "Epoch [46/200], Loss: 0.0553\n",
      "Epoch [47/200], Loss: 0.0510\n",
      "Epoch [48/200], Loss: 0.0509\n",
      "Epoch [49/200], Loss: 0.0646\n",
      "Epoch [50/200], Loss: 0.0524\n",
      "Epoch [51/200], Loss: 0.0454\n",
      "Epoch [52/200], Loss: 0.0447\n",
      "Epoch [53/200], Loss: 0.0488\n",
      "Epoch [54/200], Loss: 0.0471\n",
      "Epoch [55/200], Loss: 0.0482\n",
      "Epoch [56/200], Loss: 0.0463\n",
      "Epoch [57/200], Loss: 0.0417\n",
      "Epoch [58/200], Loss: 0.0368\n",
      "Epoch [59/200], Loss: 0.0399\n",
      "Epoch [60/200], Loss: 0.0399\n",
      "Epoch [61/200], Loss: 0.0435\n",
      "Epoch [62/200], Loss: 0.0482\n",
      "Epoch [63/200], Loss: 0.0447\n",
      "Epoch [64/200], Loss: 0.0679\n",
      "Epoch [65/200], Loss: 0.1504\n",
      "Epoch [66/200], Loss: 0.7876\n",
      "Epoch [67/200], Loss: 1.1486\n",
      "Epoch [68/200], Loss: 0.7344\n",
      "Epoch [69/200], Loss: 0.4208\n",
      "Epoch [70/200], Loss: 0.3449\n",
      "Epoch [71/200], Loss: 0.2980\n",
      "Epoch [72/200], Loss: 0.2461\n",
      "Epoch [73/200], Loss: 0.1311\n",
      "Epoch [74/200], Loss: 0.0983\n",
      "Epoch [75/200], Loss: 0.1146\n",
      "Epoch [76/200], Loss: 0.0805\n",
      "Epoch [77/200], Loss: 0.0691\n",
      "Epoch [78/200], Loss: 0.0790\n",
      "Epoch [79/200], Loss: 0.0590\n",
      "Epoch [80/200], Loss: 0.0546\n",
      "Epoch [81/200], Loss: 0.0674\n",
      "Epoch [82/200], Loss: 0.0549\n",
      "Epoch [83/200], Loss: 0.0670\n",
      "Epoch [84/200], Loss: 0.0469\n",
      "Epoch [85/200], Loss: 0.0462\n",
      "Epoch [86/200], Loss: 0.0457\n",
      "Epoch [87/200], Loss: 0.0481\n",
      "Epoch [88/200], Loss: 0.0381\n",
      "Epoch [89/200], Loss: 0.0360\n",
      "Epoch [90/200], Loss: 0.0468\n",
      "Epoch [91/200], Loss: 0.0466\n",
      "Epoch [92/200], Loss: 0.0472\n",
      "Epoch [93/200], Loss: 0.0396\n",
      "Epoch [94/200], Loss: 0.0334\n",
      "Epoch [95/200], Loss: 0.0348\n",
      "Epoch [96/200], Loss: 0.0389\n",
      "Epoch [97/200], Loss: 0.0405\n",
      "Epoch [98/200], Loss: 0.0422\n",
      "Epoch [99/200], Loss: 0.0612\n",
      "Epoch [100/200], Loss: 0.0841\n",
      "Epoch [101/200], Loss: 0.0901\n",
      "Epoch [102/200], Loss: 0.1148\n",
      "Epoch [103/200], Loss: 0.4035\n",
      "Epoch [104/200], Loss: 0.9731\n",
      "Epoch [105/200], Loss: 0.9787\n",
      "Epoch [106/200], Loss: 0.6193\n",
      "Epoch [107/200], Loss: 0.3577\n",
      "Epoch [108/200], Loss: 0.2301\n",
      "Epoch [109/200], Loss: 0.1456\n",
      "Epoch [110/200], Loss: 0.1419\n",
      "Epoch [111/200], Loss: 0.1242\n",
      "Epoch [112/200], Loss: 0.1453\n",
      "Epoch [113/200], Loss: 0.1283\n",
      "Epoch [114/200], Loss: 0.0994\n",
      "Epoch [115/200], Loss: 0.0791\n",
      "Epoch [116/200], Loss: 0.0744\n",
      "Epoch [117/200], Loss: 0.0505\n",
      "Epoch [118/200], Loss: 0.0417\n",
      "Epoch [119/200], Loss: 0.0364\n",
      "Epoch [120/200], Loss: 0.0362\n",
      "Epoch [121/200], Loss: 0.0374\n",
      "Epoch [122/200], Loss: 0.0359\n",
      "Epoch [123/200], Loss: 0.0384\n",
      "Epoch [124/200], Loss: 0.0331\n",
      "Epoch [125/200], Loss: 0.0327\n",
      "Epoch [126/200], Loss: 0.0314\n",
      "Epoch [127/200], Loss: 0.0308\n",
      "Epoch [128/200], Loss: 0.0295\n",
      "Epoch [129/200], Loss: 0.0305\n",
      "Epoch [130/200], Loss: 0.0292\n",
      "Epoch [131/200], Loss: 0.0287\n",
      "Epoch [132/200], Loss: 0.0283\n",
      "Epoch [133/200], Loss: 0.0321\n",
      "Epoch [134/200], Loss: 0.0297\n",
      "Epoch [135/200], Loss: 0.0282\n",
      "Epoch [136/200], Loss: 0.0310\n",
      "Epoch [137/200], Loss: 0.0326\n",
      "Epoch [138/200], Loss: 0.0295\n",
      "Epoch [139/200], Loss: 0.0312\n",
      "Epoch [140/200], Loss: 0.0279\n",
      "Epoch [141/200], Loss: 0.0327\n",
      "Epoch [142/200], Loss: 0.0348\n",
      "Epoch [143/200], Loss: 0.0502\n",
      "Epoch [144/200], Loss: 0.1071\n",
      "Epoch [145/200], Loss: 0.5406\n",
      "Epoch [146/200], Loss: 1.3320\n",
      "Epoch [147/200], Loss: 1.3036\n",
      "Epoch [148/200], Loss: 0.7151\n",
      "Epoch [149/200], Loss: 0.3639\n",
      "Epoch [150/200], Loss: 0.1694\n",
      "Epoch [151/200], Loss: 0.1337\n",
      "Epoch [152/200], Loss: 0.1290\n",
      "Epoch [153/200], Loss: 0.1002\n",
      "Epoch [154/200], Loss: 0.0819\n",
      "Epoch [155/200], Loss: 0.0885\n",
      "Epoch [156/200], Loss: 0.0631\n",
      "Epoch [157/200], Loss: 0.0589\n",
      "Epoch [158/200], Loss: 0.0433\n",
      "Epoch [159/200], Loss: 0.0410\n",
      "Epoch [160/200], Loss: 0.0476\n",
      "Epoch [161/200], Loss: 0.0457\n",
      "Epoch [162/200], Loss: 0.0300\n",
      "Epoch [163/200], Loss: 0.0338\n",
      "Epoch [164/200], Loss: 0.0374\n",
      "Epoch [165/200], Loss: 0.0305\n",
      "Epoch [166/200], Loss: 0.0298\n",
      "Epoch [167/200], Loss: 0.0307\n",
      "Epoch [168/200], Loss: 0.0300\n",
      "Epoch [169/200], Loss: 0.0285\n",
      "Epoch [170/200], Loss: 0.0280\n",
      "Epoch [171/200], Loss: 0.0285\n",
      "Epoch [172/200], Loss: 0.0281\n",
      "Epoch [173/200], Loss: 0.0289\n",
      "Epoch [174/200], Loss: 0.0288\n",
      "Epoch [175/200], Loss: 0.0314\n",
      "Epoch [176/200], Loss: 0.0294\n",
      "Epoch [177/200], Loss: 0.0251\n",
      "Epoch [178/200], Loss: 0.0254\n",
      "Epoch [179/200], Loss: 0.0262\n",
      "Epoch [180/200], Loss: 0.0292\n",
      "Epoch [181/200], Loss: 0.0293\n",
      "Epoch [182/200], Loss: 0.0256\n",
      "Epoch [183/200], Loss: 0.0330\n",
      "Epoch [184/200], Loss: 0.0311\n",
      "Epoch [185/200], Loss: 0.0281\n",
      "Epoch [186/200], Loss: 0.0259\n",
      "Epoch [187/200], Loss: 0.0316\n",
      "Epoch [188/200], Loss: 0.0286\n",
      "Epoch [189/200], Loss: 0.0265\n",
      "Epoch [190/200], Loss: 0.0257\n",
      "Epoch [191/200], Loss: 0.0258\n",
      "Epoch [192/200], Loss: 0.0339\n",
      "Epoch [193/200], Loss: 0.0333\n",
      "Epoch [194/200], Loss: 0.0290\n",
      "Epoch [195/200], Loss: 0.0249\n",
      "Epoch [196/200], Loss: 0.0322\n",
      "Epoch [197/200], Loss: 0.0697\n",
      "Epoch [198/200], Loss: 0.3246\n",
      "Epoch [199/200], Loss: 1.9527\n",
      "Epoch [200/200], Loss: 2.1916\n"
     ]
    }
   ],
   "execution_count": 313
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-27T07:56:39.288452Z",
     "start_time": "2025-05-27T07:56:39.278265Z"
    }
   },
   "cell_type": "code",
   "source": "model",
   "id": "a8f839a2c50b5316",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GPTModel(\n",
       "  (embedding): Embedding(233, 128)\n",
       "  (pos_encoder): PositionalEncoding()\n",
       "  (decoder): Decoder(\n",
       "    (layers): ModuleList(\n",
       "      (0-2): 3 x DecoderLayer(\n",
       "        (mha): MultiheadAttention(\n",
       "          (out_proj): NonDynamicallyQuantizableLinear(in_features=128, out_features=128, bias=True)\n",
       "        )\n",
       "        (feed_forward): FeedForward(\n",
       "          (fc1): Linear(in_features=128, out_features=512, bias=True)\n",
       "          (fc2): Linear(in_features=512, out_features=128, bias=True)\n",
       "        )\n",
       "        (layer_norm1): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "        (layer_norm2): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (linear): Linear(in_features=128, out_features=233, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 320,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 320
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-27T07:55:08.465699Z",
     "start_time": "2025-05-27T07:55:08.417641Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 推理函数\n",
    "def generate(model, sentence, vocab, device):\n",
    "    model.eval()\n",
    "    tokens = list(sentence)\n",
    "    src = torch.tensor([vocab[token] for token in tokens]).unsqueeze(0).to(device)\n",
    "    for _ in range(MAX_SEQ_LENGTH):\n",
    "\n",
    "        # 推理时，tgt是不断变化的，所以每次要生成新的掩码\n",
    "        batch_size, seq_len = src.size(0), src.size(1)\n",
    "        mask = torch.tril(torch.ones(seq_len, seq_len)).to(DEVICE)\n",
    "\n",
    "        output = model(\n",
    "            src,\n",
    "            mask=mask\n",
    "        )\n",
    "\n",
    "        # 取第二维的最后一个，第二维是sql_len，也就是取最后一个token对应的预测的结果\n",
    "        prob = output[:, -1, :]\n",
    "        # print(output.shape)\n",
    "        # print(prob.shape)\n",
    "        next_token = prob.argmax().item()\n",
    "        src = torch.cat([src, torch.tensor([[next_token]]).to(device)], dim=1)\n",
    "        # print(src)\n",
    "\n",
    "    translated = [idx_to_tgt.get(idx, '<unk>') for idx in src.squeeze().tolist()]\n",
    "    return ''.join(translated)\n",
    "\n",
    "\n",
    "# 测试翻译\n",
    "sentence = \"陈老五送进饭来\"\n",
    "print(generate(model, sentence, vocab, DEVICE))"
   ],
   "id": "fed9769a193d5a54",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "陈老五送进饭来老五，也在五，也来，\n"
     ]
    }
   ],
   "execution_count": 319
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
